package io.kwdb.sdk;

import static org.junit.jupiter.api.Assertions.*;

import io.kwdb.sdk.internal.RetryPolicy;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.jupiter.api.Test;

/**
 * Test cases for RetryPolicy component. Validates retry logic, exponential backoff, and error
 * classification.
 */
public class RetryPolicyTest {

  @Test
  public void testSuccessOnFirstAttempt() {
    RetryPolicy policy = new RetryPolicy(true, 3, 100);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              attempts.incrementAndGet();
              return Status.ok();
            });

    assertTrue(result.isOk());
    assertEquals(1, attempts.get());
  }

  @Test
  public void testRetryOnTransientError() {
    RetryPolicy policy = new RetryPolicy(true, 3, 50);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              int attempt = attempts.incrementAndGet();
              if (attempt < 3) {
                // Fail with transient error first 2 times
                return Status.unavailable("Network error");
              }
              return Status.ok();
            });

    assertTrue(result.isOk());
    assertEquals(3, attempts.get());
  }

  @Test
  public void testNoRetryOnNonTransientError() {
    RetryPolicy policy = new RetryPolicy(true, 3, 100);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              attempts.incrementAndGet();
              // Invalid argument is not transient
              return Status.invalidArgument("Bad request");
            });

    assertFalse(result.isOk());
    assertEquals(Status.Code.INVALID_ARGUMENT, result.getCode());
    assertEquals(1, attempts.get()); // Should not retry
  }

  @Test
  public void testMaxRetriesExceeded() {
    RetryPolicy policy = new RetryPolicy(true, 2, 50);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              attempts.incrementAndGet();
              // Always fail with transient error
              return Status.unavailable("Network down");
            });

    assertFalse(result.isOk());
    assertEquals(Status.Code.UNAVAILABLE, result.getCode());
    assertEquals(3, attempts.get()); // Initial + 2 retries
  }

  @Test
  public void testRetryDisabled() {
    RetryPolicy policy = new RetryPolicy(false, 3, 100);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              attempts.incrementAndGet();
              return Status.unavailable("Should not retry");
            });

    assertFalse(result.isOk());
    assertEquals(1, attempts.get()); // No retries when disabled
  }

  @Test
  public void testExponentialBackoff() {
    RetryPolicy policy = new RetryPolicy(true, 3, 100);
    AtomicInteger attempts = new AtomicInteger(0);
    long startTime = System.currentTimeMillis();

    Status result =
        policy.execute(
            () -> {
              attempts.incrementAndGet();
              return Status.unavailable("Network error");
            });

    long elapsed = System.currentTimeMillis() - startTime;

    // With exponential backoff, should have delays between retries
    // Base: 100ms, then ~200ms, then ~400ms (with jitter)
    // Total should be at least 500ms (allowing for jitter)
    assertTrue(elapsed >= 400, "Should have exponential backoff delays");
    assertEquals(4, attempts.get()); // Initial + 3 retries
  }

  @Test
  public void testRetryOnDeadlineExceeded() {
    RetryPolicy policy = new RetryPolicy(true, 2, 50);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              int attempt = attempts.incrementAndGet();
              if (attempt == 1) {
                return Status.deadlineExceeded("Timeout");
              }
              return Status.ok();
            });

    assertTrue(result.isOk());
    assertEquals(2, attempts.get());
  }

  @Test
  public void testRetryOnResourceExhausted() {
    RetryPolicy policy = new RetryPolicy(true, 2, 50);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              int attempt = attempts.incrementAndGet();
              if (attempt == 1) {
                return Status.resourceExhausted("Queue full");
              }
              return Status.ok();
            });

    assertTrue(result.isOk());
    assertEquals(2, attempts.get());
  }

  @Test
  public void testRetryOnAborted() {
    RetryPolicy policy = new RetryPolicy(true, 2, 50);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              int attempt = attempts.incrementAndGet();
              if (attempt == 1) {
                return Status.aborted("Transaction conflict");
              }
              return Status.ok();
            });

    assertTrue(result.isOk());
    assertEquals(2, attempts.get());
  }

  @Test
  public void testExceptionHandling() {
    RetryPolicy policy = new RetryPolicy(true, 2, 50);
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              int attempt = attempts.incrementAndGet();
              if (attempt < 3) {
                throw new RuntimeException("Connection failed");
              }
              return Status.ok();
            });

    assertTrue(result.isOk());
    assertEquals(3, attempts.get());
  }

  @Test
  public void testInterruptedExecution() throws InterruptedException {
    RetryPolicy policy = new RetryPolicy(true, 3, 1000);

    Thread testThread =
        new Thread(
            () -> {
              Status result =
                  policy.execute(
                      () -> {
                        // Simulate long operation
                        Thread.sleep(5000);
                        return Status.ok();
                      });

              assertFalse(result.isOk());
              assertEquals(Status.Code.FAILED_PRECONDITION, result.getCode());
            });

    testThread.start();
    Thread.sleep(100); // Let it start
    testThread.interrupt();
    testThread.join(1000);
    assertFalse(testThread.isAlive());
  }

  @Test
  public void testDefaultPolicy() {
    RetryPolicy policy = RetryPolicy.defaultPolicy();
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              int attempt = attempts.incrementAndGet();
              if (attempt == 1) {
                return Status.unavailable("Temporary failure");
              }
              return Status.ok();
            });

    assertTrue(result.isOk());
    assertEquals(2, attempts.get());
  }

  @Test
  public void testNoRetryPolicy() {
    RetryPolicy policy = RetryPolicy.noRetry();
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              attempts.incrementAndGet();
              return Status.unavailable("Should not retry");
            });

    assertFalse(result.isOk());
    assertEquals(1, attempts.get());
  }

  @Test
  public void testMaxRetriesLimit() {
    // Test that max retries is capped at 10
    RetryPolicy policy = new RetryPolicy(true, 100, 10); // Request 100 retries
    AtomicInteger attempts = new AtomicInteger(0);

    Status result =
        policy.execute(
            () -> {
              attempts.incrementAndGet();
              return Status.unavailable("Always fail");
            });

    assertFalse(result.isOk());
    // Should be capped at 10 retries + 1 initial = 11 total
    assertTrue(attempts.get() <= 11, "Retries should be capped at 10");
  }

  @Test
  public void testJitterInBackoff() {
    RetryPolicy policy = new RetryPolicy(true, 5, 100);

    // Run multiple times to test jitter variation
    long[] delays = new long[3];

    for (int i = 0; i < delays.length; i++) {
      AtomicInteger attempts = new AtomicInteger(0);
      long start = System.currentTimeMillis();

      policy.execute(
          () -> {
            if (attempts.incrementAndGet() < 2) {
              return Status.unavailable("Fail once");
            }
            return Status.ok();
          });

      delays[i] = System.currentTimeMillis() - start;
    }

    // Delays should vary due to jitter
    boolean hasVariation = false;
    for (int i = 1; i < delays.length; i++) {
      if (Math.abs(delays[i] - delays[0]) > 10) { // Allow 10ms difference
        hasVariation = true;
        break;
      }
    }
    assertTrue(hasVariation, "Backoff should have jitter variation");
  }
}
